#!/usr/bin/env python
"""
:mod:`pynchon.bio.alg.dirichlet`
================================

Algorithms on dirichlet mixtures.
"""
from itertools import izip
from pynchon.core.symbol import Symbol
import numpy as np
import sympy as sy

def _validate(count1, count2, dirichlet):
    # preliminary check dirichlet
    try:
        symbols, weights, alpha_vals = dirichlet
    except (ValueError, TypeError):
        raise ValueError('Invalid Dirichlet mixture.')
    # count2 might be None
    count2 = count2 if count2 is not None else count1
    # check if all have the same shape
    try:
        assert count1.shape == count2.shape
        assert len(alpha_vals[0]) == len(count1)
        assert len(alpha_vals) == len(weights)

    except AssertionError:
        raise ValueError('Input array shape mismatch.')
    # check if counts have the same alphabet as the dirichlet
    for cnt in (count1, count2):
        if isinstance(cnt, Symbol):
            if not cnt.metadata['alphabet'] == symbols:
                raise ValueError('Incompatible symbols or alphabet order.')        
    return weights, alpha_vals

def p_count1(count, dirichlet):
    """
    Calculates the probability of indpendently sampling a count vector
    given a Dirichlet mixture parameterization of background distributions.
    Returns a an unevaluated SymPy object.
    
    Arguments:
    
        - count (``numpy.ndarray``) A ``Freq`` instance or ``numpy.ndarray`` or
          sequence of counts.  
        - dirichlet (``sequence``) A tuple of (symbols, weights, components), 
          where "symbols" is a sequence of symbols, "weights" is a sequence of 
          weights for each mixture and "components" is a rank-2 array of alphas,
          where each component is a row. See: 
          ``pynchon.bio.data.dirichlet.dist_20comp`` for an example.
    """
    weights, alpha_vals = _validate(count, None, dirichlet)
    Pn = 0.
    for weight, a_val in izip(weights, alpha_vals):
        # substitute
        Za = beta_function(a_val)
        Zna = beta_function(count + a_val)
        Mn = multinomial_coeff(count.tolist())
        eqr = Zna * Mn / Za
        Pn += weight * eqr
    return Pn

def p_count2(count1, count2, dirichlet):
    """
    Calculates the probability of indpendently sampling two count vectors 
    given a Dirichlet mixture parameterization of background distributions.
    Returns a unevaluated SymPy object.
    
    Arguments:
    
        - count1 (``numpy.ndarray``) A ``Freq`` instance or ``numpy.ndarray`` or
          sequence of counts.
        - count2 (``numpy.ndarray``) A ``Freq`` instance or ``numpy.ndarray`` or
          sequence of counts.          
        - dirichlet (``sequence``) A tuples of (symbols, weights, components), 
          where "symbols" is a sequence of symbols, "weights" is a sequence of 
          weights for each mixture and "components" is a rank-2 array of alphas,
          where each component is a row. See: 
          ``pynchon.bio.data.dirichlet.dist_20comp`` for an example.
    """
    weights, alpha_vals = _validate(count1, count2, dirichlet)
    Pn1n2 = 0.
    for weight, a_val in izip(weights, alpha_vals):
        Za = beta_function(a_val)
        Zn1n2a = beta_function(count1 + count2 + a_val)
        Mn1 = multinomial_coeff(count1.tolist())
        Mn2 = multinomial_coeff(count2.tolist())
        eqr = (Zn1n2a * Mn1 * Mn2) / Za
        Pn1n2 += weight * eqr
    return Pn1n2

def s_count2(count1, count2, dirichlet):
    """
    Calculates a bit-score between two count vectors, given a dirichlet mixture
    prior. Returns an unevaluated SymPy object.
    
    Arguments:
    
        - count1 (``numpy.ndarray``) A ``Freq`` instance or ``numpy.ndarray`` or
          sequence of counts.
        - count2 (``numpy.ndarray``) A ``Freq`` instance or ``numpy.ndarray`` or
          sequence of counts.          
        - dirichlet (``sequence``) A tuples of (symbols, weights, components), 
          where "symbols" is a sequence of symbols, "weights" is a sequence of 
          weights for each mixture and "components" is a rank-2 array of alphas,
          where each component is a row. See: 
          ``pynchon.bio.data.dirichlet.dist_20comp`` for an example.    
    """
    Pn1 = p_count1(count1, dirichlet)
    Pn2 = p_count1(count2, dirichlet)
    Pn1n2 = p_count2(count1, count2, dirichlet)
    Sn1n2 = sy.log((Pn1n2 / (Pn1*Pn2)), 2) # return result in bits
    return Sn1n2

def p_count1_bg(count, background):
    """
    Calculates the probability of independently sampling a count vector from a 
    background distribution.
    
    Arguments
    
        - count (``numpy.ndarray``) A ``Freq`` instance or ``numpy.ndarray`` of
          rank-1 or sequence of counts. 
        - background (``numpy.ndarray``) A ``Freq`` instance or 
          ``numpy.ndarray`` with frequency vectors along the last axis. This 
          typically is a rank-1 array.
    """
    positives = (count != 0.) & (background != 0.)  
    pos_cts = count[positives]
    pos_bgs = background[positives]
    # to scalars
    Mn = multinomial_coeff(pos_cts)
    th = np.prod(pos_bgs**pos_cts)
    P = Mn * th
    return P

def p_symbols(count, dirichlet):
    """
    Calculates the probability of each symbol in the alphabet given a 
    dirichlet mixture and a observed count vector. Returns a sequence of 
    SymPy evaluated Reals.
    
    Arguments:

        - count (``numpy.ndarray``) A ``Freq`` instance or ``numpy.ndarray`` or
          sequence of counts.  
        - dirichlet (``sequence``) A tuple of (symbols, weights, components), 
          where "symbols" is a sequence of letters, "weights" is a sequence of 
          weights for each mixture and "components" is a rank-2 array of alpha
          values, where each component is a row. See: 
          ``pynchon.bio.data.dirichlet.dist_20comp`` for an example.    
    """
    symbols = dirichlet[0]
    weights, alphas = _validate(count, None, dirichlet)   
    Pn = p_count1(count, dirichlet)
    n_sum = np.sum(count)
    na = []
    Pa = []
    Ps = []
    for j in xrange(len(alphas)):
        weight = weights[j:j+1]
        alpha = alphas[j:j+1]
        paj = p_count1(count, (symbols, weight, alpha)) / Pn
        Pa.append(paj)
        na.append(np.sum(alpha) + n_sum)
    for i in xrange(len(symbols)):
        psi = 0.   
        ni = count[i]
        for j in xrange(len(alphas)):
            psij = Pa[j] * ((ni + alphas[j][i]) / na[j])
            psi += psij
        Ps.append(psi.evalf())
    return Ps

def multinomial_coeff(count):
    """
    Returns the multinomial coefficient for a given count vector. Returns a 
    SymPy object.
    
    Arguments:
    
        - count (``sequence`` or ``numpy.ndarray``) A vector of integer counts.
    """
    ln = len(count)
    n = [sy.S('n%s' % i) for i in range(ln)]
    Mn = sy.factorial(sy.Add(*n))/sy.Mul(*sy.factorial(n))
    return Mn.subs(zip(n, count))

def beta_function(vector):
    """
    Calculates the value of the Beta function for a non-negative vector. Returns
    a SymPy object. 
    
    Arguments:
    
        - vector (``sequence`` or ``numpy.ndarray``) A vector of floating point
          numbers. 
    """
    ln = len(vector)
    a = [sy.S('a%s' % i) for i in range(ln)]
    Za = sy.Mul(*sy.gamma(a))/sy.gamma(sy.Add(*a))
    return Za.subs(zip(a, vector))
